import os, glob, re
import torch
from typing import Any

from torchvision.transforms import Compose
from sklearn.model_selection import train_test_split
from ecgcmr.utils.ecg_parser import read_ecg
from ecgcmr.signal.augmentations.CropResizing import CropResizing
from ecgcmr.signal.augmentations.BaselineCorrection import BaselineCorrection
from ecgcmr.signal.augmentations.Normalize import Normalize


class ECGDataset:
    def __init__(self, data_path, split='train', transform=True, args=None, val_size=0.2, random_state=42) -> None:

        all_samples = self._load_samples(data_path)

        train_samples, val_samples = train_test_split(all_samples,
                                                      test_size=val_size,
                                                      random_state=random_state)
        self.split = split
        
        self.samples = train_samples if self.split == 'train' else val_samples

        self.transform = self._compose() if transform else None
        
        self.args = args

    def _load_samples(self, data_path):
        pattern = re.compile(r"\d+_20205_\d_\d\.xml$") # take only rest ecg
        return [file for file in glob.iglob(os.path.join(data_path, '*.xml')) if pattern.match(os.path.basename(file))]
    
    def _compose(self):
        steps = []
        steps.append(BaselineCorrection())
        steps.append(Normalize())
        steps.append(CropResizing(fixed_crop_len=self.args.input_size[-1], start_idx=0, resize=False))

        return Compose(steps)

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx):

        raw_ecg_path = self.samples[idx]

        raw_ecg, md = read_ecg(raw_ecg_path)
        
        if self.transform:
            ecg = self.transform(raw_ecg)
        else:
            ecg = raw_ecg

        ecg = torch.tensor(ecg, dtype=torch.double).unsqueeze(dim=0)

        return ecg, md

